require 'rubygems'
require 'mechanize'

module StadigTech
  module Edgarr
    COMPANY_SEARCH_URL = 'http://sec.gov/edgar/searchedgar/companysearch.html'
    CIK_LOOKUP_URL = 'http://sec.gov/edgar/searchedgar/cik.htm'
    class Company
      def self.lookup_cik(company)
        agent = WWW::Mechanize.new
        page = agent.get(CIK_LOOKUP_URL)
        cik_form = page.forms.first
        cik_form.company = company
        page = agent.submit(cik_form)
        page.links[0...-1].collect do |l|
          [l.text, l.node.parent.children[l.node.parent.children.index(l.node)+1].content.to_s.strip]
        end
      end
      
      def self.find_by_name(name)
        agent = WWW::Mechanize.new do |m|
          m.user_agent = WWW::Mechanize::AGENT_ALIASES['Windows IE 6']
        end
        page = agent.get(COMPANY_SEARCH_URL)
        form = page.forms.first
        form.company = name
        page = agent.submit(form)
        if page.title == 'Matching Company Name Results'
          companies = parse_multi_results(page, [])
          while !page.forms.empty? && (page.root.search('//form/input[@type=button]').size == 2 || page.root.search('//form/input[@type=button]')[0].attributes['value'][0, 4] == 'Next')
            if page.root.search('//form/input[@type=button]').size < 2
              url = 'http://sec.gov' + page.root.search('//form/input:eq(0)').first.attributes["onclick"][17...-1]
            else
              url = 'http://sec.gov' + page.root.search('//form/input:eq(1)').first.attributes["onclick"][17...-1]
            end
            page = agent.get(url)
            companies = parse_multi_results(page, companies)
          end
          companies
        elsif page.title[0, 19] == 'Company Information'
          [Company.new(page.links[5].text)]
        else
          raise "Unknown response"
        end
      end
      
      def self.find_by_ticker(ticker)
        agent = WWW::Mechanize.new
        page = agent.get(COMPANY_SEARCH_URL)
        form = page.forms.first
        form.cik = ticker
        page = agent.submit(form)
        Company.new(page.forms.first)
      end
      
      attr_reader :cik, :name, :state, :sic
      
      def initialize(*args)
        if block_given?
          yield self
        elsif args.length != 1
          raise 'Wrong number of arguments'
        elsif args[0].is_a?(Hash)
          @cik = args[0][:cik].to_i
          @name = args[0][:name].to_s
          @state = args[0][:state].to_s
          @sic = args[0][:sic].to_i
        else
          @cik = args[0].to_i
        end
        @former_names = []
      end
      
      def cik=(val)
        @cik = val.to_i
      end
      
      def name=(val)
        @name = val.to_s
      end
      
      def state=(val)
        @state = val.to_s
      end
      
      def sic=(val)
        @sic = val.to_i
      end
      
      def former_name?
        !@former_names.empty?
      end
      
      def add_former_name(date, name)
        found = nil
        @former_names.each_index do |i|
          if date <= @former_names[i][0]
            found = i
            break
          end
        end
        
        if found
          @former_names.insert(found, [date, name])
        else
          @former_names << [date, name]
        end
      end
      
      def former_name(date)
        @former_names.each do |n|
          if n[0] <= date
            return n[1]
          end
        end
      end
      
      def each_former_name
        @former_names.each do |n|
          yield n[0], n[1]
        end
      end
      
      private
      def self.parse_multi_results(page, companies)
        curr = nil
        page.root.search('html/body/div/table/tr').inject(companies) do |result, tr|
          if !tr.search('th').empty?
            result
          elsif tr.search('td:eq(0)/a').text =~ /^\d+$/
            curr = Company.new do |c|
              c.cik = tr.search('td:eq(0)/a:eq(0)').text
              c.name = tr.search('td:eq(1)').text
              c.state = tr.search('td:eq(2)/a:eq(0)').text
            end
            result << curr
          else
            if (tr/'td:eq(0)').text == 'SIC:'
              curr.sic = tr.search('td:eq(1)/a:eq(0)').text if curr
            elsif (tr/'td:eq(1)').text =~ /^([^(]+)\(filings through.(\d{4})-(\d{2})-(\d{2})\).*$/
             (tr/'td:eq(1)').text.split("\n").each do |s|
                if s =~ /^([^(]+)\(filings through.(\d{4})-(\d{2})-(\d{2})\).*$/
                  curr.add_former_name($1, Date.civil($2.to_i, $3.to_i, $4.to_i))
                end
              end
            else
              raise "Something odd: #{tr.to_s}"
            end
            result
          end
        end
      end
    end
  end
end
